from langchain_milvus import Milvus
from langchain_ollama import OllamaEmbeddings  # 可替换为其他嵌入模型
import doc_split

# 初始化Milvus向量存储（类似ORM的模型定义）
embeddings = OllamaEmbeddings(
    model='bge-m3'
)
vector_db = Milvus(
    embedding_function=embeddings,  # 向量生成器
    connection_args={"uri": "http://123.60.131.133:19530"},
    collection_name="langchain_demo",
    auto_id=True,
    drop_old=True  # 首次运行可删除旧集合
)

# 插入数据（文本会自动转为向量）
texts = doc_split.split_file("c:/ai/doc/消费者权益保护法.txt")
print("文字匹配------------")
for text in texts:
    if '第七章 第五十' in text:
        print(text)
vector_db.add_texts(texts=texts)

# 检索相似文本（自动生成查询向量并检索）
query = "买到假货可以申请多少钱的赔偿"
docs = vector_db.similarity_search(query, k=3)
# print(docs[0].page_content)  # 输出匹配的文本
print("模型匹配 ===============")
for doc in docs:
    print(f'{len(doc.page_content)} \t {doc.page_content}')